#include "jit_compiler.h"
#include "../memory/memory_manager.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

// =============================================================================
// Function Management
// =============================================================================

jit_function_t* jit_function_create(jit_compiler_t* compiler, const char* name, 
                                   jit_type_t return_type, jit_type_t* param_types, 
                                   uint32_t param_count) {
    if (!compiler || !name || param_count > JIT_MAX_PARAMETERS) return NULL;
    
    jit_function_t* function = mm_alloc(compiler->memory_mgr, sizeof(jit_function_t), 
                                       MEM_TYPE_QUERY_BUFFER);
    if (!function) return NULL;
    
    memset(function, 0, sizeof(jit_function_t));
    
    // Initialize basic fields
    strncpy(function->name, name, sizeof(function->name) - 1);
    function->return_type = return_type;
    function->param_count = param_count;
    
    // Copy parameter types
    if (param_count > 0 && param_types) {
        function->param_types = mm_alloc(compiler->memory_mgr, 
                                       param_count * sizeof(jit_type_t),
                                       MEM_TYPE_QUERY_BUFFER);
        if (!function->param_types) {
            mm_free(compiler->memory_mgr, function);
            return NULL;
        }
        memcpy(function->param_types, param_types, param_count * sizeof(jit_type_t));
    }
    
    // Initialize values array
    function->values = mm_alloc(compiler->memory_mgr, 
                               256 * sizeof(jit_value_t*), // Initial capacity
                               MEM_TYPE_QUERY_BUFFER);
    if (!function->values) {
        if (function->param_types) mm_free(compiler->memory_mgr, function->param_types);
        mm_free(compiler->memory_mgr, function);
        return NULL;
    }
    memset(function->values, 0, 256 * sizeof(jit_value_t*));
    
    // Create entry block
    function->entry_block = jit_block_create(function);
    if (!function->entry_block) {
        mm_free(compiler->memory_mgr, function->values);
        if (function->param_types) mm_free(compiler->memory_mgr, function->param_types);
        mm_free(compiler->memory_mgr, function);
        return NULL;
    }
    function->blocks = function->entry_block;
    function->block_count = 1;
    
    // Initialize synchronization
    if (pthread_mutex_init(&function->mutex, NULL) != 0) {
        // Continue without mutex (degraded performance but still functional)
    }
    
    // Set initial optimization level
    function->optimization_level = compiler->optimization_level;
    
    return function;
}

void jit_function_destroy(jit_function_t* function) {
    if (!function) return;
    
    // Clean up basic blocks
    jit_basic_block_t* block = function->blocks;
    while (block) {
        jit_basic_block_t* next = block->next;
        
        // Clean up instructions in block
        jit_instruction_t* inst = block->first_inst;
        while (inst) {
            jit_instruction_t* next_inst = inst->next;
            if (inst->machine_code) {
                free(inst->machine_code);
            }
            free(inst);
            inst = next_inst;
        }
        
        // Clean up block machine code
        if (block->machine_code) {
            free(block->machine_code);
        }
        
        // Clean up predecessors array
        if (block->predecessors) {
            free(block->predecessors);
        }
        
        free(block);
        block = next;
    }
    
    // Clean up values
    if (function->values) {
        for (uint32_t i = 0; i < function->value_count; i++) {
            if (function->values[i]) {
                if (function->values[i]->type == JIT_TYPE_STRING && 
                    function->values[i]->data.string_val) {
                    free(function->values[i]->data.string_val);
                }
                free(function->values[i]);
            }
        }
        free(function->values);
    }
    
    // Clean up parameter types
    if (function->param_types) {
        free(function->param_types);
    }
    
    // Clean up machine code
    if (function->machine_code) {
        free(function->machine_code);
    }
    
    // Destroy mutex
    pthread_mutex_destroy(&function->mutex);
    
    free(function);
}

int jit_function_compile(jit_function_t* function) {
    if (!function || function->is_compiled) return -1;
    
    pthread_mutex_lock(&function->mutex);
    
    // Validate function
    if (jit_function_validate(function) != 0) {
        pthread_mutex_unlock(&function->mutex);
        return -1;
    }
    
    // Run optimization passes
    if (jit_optimize_function(function, function->optimization_level) != 0) {
        pthread_mutex_unlock(&function->mutex);
        return -1;
    }
    
    // Generate machine code
    if (jit_generate_machine_code(function) != 0) {
        pthread_mutex_unlock(&function->mutex);
        return -1;
    }
    
    function->is_compiled = true;
    pthread_mutex_unlock(&function->mutex);
    
    return 0;
}

// =============================================================================
// Basic Block Management
// =============================================================================

jit_basic_block_t* jit_block_create(jit_function_t* function) {
    if (!function) return NULL;
    
    jit_basic_block_t* block = malloc(sizeof(jit_basic_block_t));
    if (!block) return NULL;
    
    memset(block, 0, sizeof(jit_basic_block_t));
    block->id = function->block_count;
    
    // Link into function
    if (function->blocks) {
        // Find last block and append
        jit_basic_block_t* last = function->blocks;
        while (last->next) {
            last = last->next;
        }
        last->next = block;
    } else {
        function->blocks = block;
    }
    function->block_count++;
    
    return block;
}

int jit_block_add_successor(jit_basic_block_t* block, jit_basic_block_t* successor) {
    if (!block || !successor) return -1;
    
    // Add successor
    if (!block->successor1) {
        block->successor1 = successor;
    } else if (!block->successor2) {
        block->successor2 = successor;
    } else {
        return -1; // Already has 2 successors
    }
    
    // Add this block as predecessor to successor
    if (successor->predecessors == NULL) {
        successor->predecessors = malloc(sizeof(jit_basic_block_t*));
        if (!successor->predecessors) return -1;
        successor->predecessors[0] = block;
        successor->predecessor_count = 1;
    } else {
        // Reallocate predecessors array
        jit_basic_block_t** new_preds = realloc(successor->predecessors,
                                              (successor->predecessor_count + 1) * sizeof(jit_basic_block_t*));
        if (!new_preds) return -1;
        successor->predecessors = new_preds;
        successor->predecessors[successor->predecessor_count] = block;
        successor->predecessor_count++;
    }
    
    return 0;
}

// =============================================================================
// Value Management
// =============================================================================

jit_value_t* jit_value_create_const_int(jit_context_t* ctx, int64_t value) {
    if (!ctx || !ctx->current_function) return NULL;
    
    jit_value_t* jit_val = malloc(sizeof(jit_value_t));
    if (!jit_val) return NULL;
    
    memset(jit_val, 0, sizeof(jit_value_t));
    jit_val->type = JIT_TYPE_INT64;
    jit_val->id = ctx->current_function->next_value_id++;
    jit_val->is_constant = true;
    jit_val->is_parameter = false;
    jit_val->data.int_val = value;
    jit_val->reg_id = -1; // Not assigned to register yet
    
    // Add to function's values array
    if (ctx->current_function->value_count < 256) { // Assume initial capacity
        ctx->current_function->values[ctx->current_function->value_count] = jit_val;
        ctx->current_function->value_count++;
    }
    
    return jit_val;
}

jit_value_t* jit_value_create_const_float(jit_context_t* ctx, double value) {
    if (!ctx || !ctx->current_function) return NULL;
    
    jit_value_t* jit_val = malloc(sizeof(jit_value_t));
    if (!jit_val) return NULL;
    
    memset(jit_val, 0, sizeof(jit_value_t));
    jit_val->type = JIT_TYPE_DOUBLE;
    jit_val->id = ctx->current_function->next_value_id++;
    jit_val->is_constant = true;
    jit_val->is_parameter = false;
    jit_val->data.float_val = value;
    jit_val->reg_id = -1;
    
    // Add to function's values array
    if (ctx->current_function->value_count < 256) {
        ctx->current_function->values[ctx->current_function->value_count] = jit_val;
        ctx->current_function->value_count++;
    }
    
    return jit_val;
}

jit_value_t* jit_value_create_const_string(jit_context_t* ctx, const char* value) {
    if (!ctx || !ctx->current_function || !value) return NULL;
    
    jit_value_t* jit_val = malloc(sizeof(jit_value_t));
    if (!jit_val) return NULL;
    
    memset(jit_val, 0, sizeof(jit_value_t));
    jit_val->type = JIT_TYPE_STRING;
    jit_val->id = ctx->current_function->next_value_id++;
    jit_val->is_constant = true;
    jit_val->is_parameter = false;
    jit_val->data.string_val = strdup(value);
    jit_val->reg_id = -1;
    
    if (!jit_val->data.string_val) {
        free(jit_val);
        return NULL;
    }
    
    // Add to function's values array
    if (ctx->current_function->value_count < 256) {
        ctx->current_function->values[ctx->current_function->value_count] = jit_val;
        ctx->current_function->value_count++;
    }
    
    return jit_val;
}

jit_value_t* jit_value_create_parameter(jit_context_t* ctx, uint32_t param_index, jit_type_t type) {
    if (!ctx || !ctx->current_function || param_index >= ctx->current_function->param_count) {
        return NULL;
    }
    
    jit_value_t* jit_val = malloc(sizeof(jit_value_t));
    if (!jit_val) return NULL;
    
    memset(jit_val, 0, sizeof(jit_value_t));
    jit_val->type = type;
    jit_val->id = ctx->current_function->next_value_id++;
    jit_val->is_constant = false;
    jit_val->is_parameter = true;
    jit_val->reg_id = -1;
    
    // Store parameter index in data (for code generation)
    jit_val->data.uint_val = param_index;
    
    // Add to function's values array
    if (ctx->current_function->value_count < 256) {
        ctx->current_function->values[ctx->current_function->value_count] = jit_val;
        ctx->current_function->value_count++;
    }
    
    return jit_val;
}

jit_value_t* jit_value_create_temp(jit_context_t* ctx, jit_type_t type) {
    if (!ctx || !ctx->current_function) return NULL;
    
    jit_value_t* jit_val = malloc(sizeof(jit_value_t));
    if (!jit_val) return NULL;
    
    memset(jit_val, 0, sizeof(jit_value_t));
    jit_val->type = type;
    jit_val->id = ctx->current_function->next_value_id++;
    jit_val->is_constant = false;
    jit_val->is_parameter = false;
    jit_val->reg_id = -1;
    
    // Add to function's values array
    if (ctx->current_function->value_count < 256) {
        ctx->current_function->values[ctx->current_function->value_count] = jit_val;
        ctx->current_function->value_count++;
    }
    
    return jit_val;
}

// =============================================================================
// Instruction Generation
// =============================================================================

jit_instruction_t* jit_emit_binary_op(jit_context_t* ctx, jit_op_t op, 
                                     jit_value_t* left, jit_value_t* right) {
    if (!ctx || !ctx->current_block || !left || !right) return NULL;
    
    jit_instruction_t* inst = malloc(sizeof(jit_instruction_t));
    if (!inst) return NULL;
    
    memset(inst, 0, sizeof(jit_instruction_t));
    inst->opcode = op;
    inst->src1 = left;
    inst->src2 = right;
    
    // Create destination value
    jit_type_t result_type = jit_type_promote(left->type, right->type);
    inst->dest = jit_value_create_temp(ctx, result_type);
    if (!inst->dest) {
        free(inst);
        return NULL;
    }
    
    // Add instruction to current block
    if (!ctx->current_block->first_inst) {
        ctx->current_block->first_inst = inst;
        ctx->current_block->last_inst = inst;
    } else {
        inst->prev = ctx->current_block->last_inst;
        ctx->current_block->last_inst->next = inst;
        ctx->current_block->last_inst = inst;
    }
    ctx->current_block->inst_count++;
    
    return inst;
}

jit_instruction_t* jit_emit_unary_op(jit_context_t* ctx, jit_op_t op, jit_value_t* operand) {
    if (!ctx || !ctx->current_block || !operand) return NULL;
    
    jit_instruction_t* inst = malloc(sizeof(jit_instruction_t));
    if (!inst) return NULL;
    
    memset(inst, 0, sizeof(jit_instruction_t));
    inst->opcode = op;
    inst->src1 = operand;
    
    // Create destination value
    inst->dest = jit_value_create_temp(ctx, operand->type);
    if (!inst->dest) {
        free(inst);
        return NULL;
    }
    
    // Add instruction to current block
    if (!ctx->current_block->first_inst) {
        ctx->current_block->first_inst = inst;
        ctx->current_block->last_inst = inst;
    } else {
        inst->prev = ctx->current_block->last_inst;
        ctx->current_block->last_inst->next = inst;
        ctx->current_block->last_inst = inst;
    }
    ctx->current_block->inst_count++;
    
    return inst;
}

jit_instruction_t* jit_emit_return(jit_context_t* ctx, jit_value_t* value) {
    if (!ctx || !ctx->current_block) return NULL;
    
    jit_instruction_t* inst = malloc(sizeof(jit_instruction_t));
    if (!inst) return NULL;
    
    memset(inst, 0, sizeof(jit_instruction_t));
    inst->opcode = JIT_OP_RETURN;
    inst->src1 = value;
    
    // Add instruction to current block
    if (!ctx->current_block->first_inst) {
        ctx->current_block->first_inst = inst;
        ctx->current_block->last_inst = inst;
    } else {
        inst->prev = ctx->current_block->last_inst;
        ctx->current_block->last_inst->next = inst;
        ctx->current_block->last_inst = inst;
    }
    ctx->current_block->inst_count++;
    
    return inst;
}

// =============================================================================
// Context Management
// =============================================================================

jit_context_t* jit_context_create(jit_compiler_t* compiler) {
    if (!compiler) return NULL;
    
    jit_context_t* ctx = mm_alloc(compiler->memory_mgr, sizeof(jit_context_t), 
                                 MEM_TYPE_TEMP_BUFFER);
    if (!ctx) return NULL;
    
    memset(ctx, 0, sizeof(jit_context_t));
    
    // Initialize value stack
    ctx->stack_capacity = 64;
    ctx->value_stack = mm_alloc(compiler->memory_mgr, 
                               ctx->stack_capacity * sizeof(jit_value_t*),
                               MEM_TYPE_TEMP_BUFFER);
    if (!ctx->value_stack) {
        mm_free(compiler->memory_mgr, ctx);
        return NULL;
    }
    
    // Initialize temporary buffer
    ctx->temp_buffer_size = 4096;
    ctx->temp_buffer = mm_alloc(compiler->memory_mgr, ctx->temp_buffer_size, 
                               MEM_TYPE_TEMP_BUFFER);
    if (!ctx->temp_buffer) {
        mm_free(compiler->memory_mgr, ctx->value_stack);
        mm_free(compiler->memory_mgr, ctx);
        return NULL;
    }
    
    return ctx;
}

void jit_context_destroy(jit_context_t* ctx) {
    if (!ctx) return;
    
    if (ctx->value_stack) {
        free(ctx->value_stack);
    }
    
    if (ctx->temp_buffer) {
        free(ctx->temp_buffer);
    }
    
    free(ctx);
}

void jit_block_set_current(jit_context_t* ctx, jit_basic_block_t* block) {
    if (ctx) {
        ctx->current_block = block;
    }
}

int jit_context_push_value(jit_context_t* ctx, jit_value_t* value) {
    if (!ctx || !value || ctx->stack_top >= ctx->stack_capacity) return -1;
    
    ctx->value_stack[ctx->stack_top++] = value;
    return 0;
}

jit_value_t* jit_context_pop_value(jit_context_t* ctx) {
    if (!ctx || ctx->stack_top == 0) return NULL;
    
    return ctx->value_stack[--ctx->stack_top];
}

// =============================================================================
// Function Validation
// =============================================================================

int jit_function_validate(const jit_function_t* function) {
    if (!function) return -1;
    
    // Check that function has at least one block
    if (!function->blocks) return -1;
    
    // Check that entry block exists
    if (!function->entry_block) return -1;
    
    // Validate each block
    jit_basic_block_t* block = function->blocks;
    while (block) {
        // Check that block has at least one instruction
        if (!block->first_inst) {
            // Empty blocks are allowed in some cases
        }
        
        // Validate instructions
        jit_instruction_t* inst = block->first_inst;
        while (inst) {
            // Check that instruction has valid opcode
            if (inst->opcode >= 64) return -1; // Assuming max 64 opcodes
            
            // Check that instruction operands are valid
            if (inst->src1 && inst->src1->type >= JIT_TYPE_UNKNOWN) return -1;
            if (inst->src2 && inst->src2->type >= JIT_TYPE_UNKNOWN) return -1;
            if (inst->dest && inst->dest->type >= JIT_TYPE_UNKNOWN) return -1;
            
            inst = inst->next;
        }
        
        block = block->next;
    }
    
    return 0;
}

// =============================================================================
// Debug and Introspection
// =============================================================================

void jit_function_print_ir(const jit_function_t* function) {
    if (!function) return;
    
    printf("Function: %s\n", function->name);
    printf("Return type: %s\n", jit_type_to_string(function->return_type));
    printf("Parameters: %u\n", function->param_count);
    for (uint32_t i = 0; i < function->param_count; i++) {
        printf("  param_%u: %s\n", i, jit_type_to_string(function->param_types[i]));
    }
    printf("Blocks: %u\n", function->block_count);
    printf("Values: %u\n", function->value_count);
    printf("\n");
    
    // Print blocks and instructions
    jit_basic_block_t* block = function->blocks;
    while (block) {
        printf("Block %u:\n", block->id);
        
        jit_instruction_t* inst = block->first_inst;
        while (inst) {
            printf("  ");
            
            // Print destination
            if (inst->dest) {
                printf("v%u = ", inst->dest->id);
            }
            
            // Print opcode
            printf("%d", (int)inst->opcode);
            
            // Print operands
            if (inst->src1) {
                printf(" v%u", inst->src1->id);
            }
            if (inst->src2) {
                printf(" v%u", inst->src2->id);
            }
            if (inst->src3) {
                printf(" v%u", inst->src3->id);
            }
            
            printf("\n");
            inst = inst->next;
        }
        
        printf("\n");
        block = block->next;
    }
}